package SparkSQL

import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}

/**
 * 弱类型
 * UDAF是自定义函数
 * 他和UDF不一样的地方是UDF是一对一输出的
 * 但是UDAF是聚合函数（类似于sum、count）
 */
object UDAFTest {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("UDAF")
      .getOrCreate()
    val df = spark.read
      .json("data/user.json")
    df.createOrReplaceTempView("user")

    //注册函数
    spark.udf.register("ageAvg",new MyAvgUDAF())

    spark.sql("select ageAvg(age) from user").show()

    spark.close()
  }

  /**
   * 自定义聚合函数,计算年龄的平均值
   * 继承：UserDefinedAggregateFunction
   */
  class MyAvgUDAF extends UserDefinedAggregateFunction {
    /**
     * 输入数据的结构，我们这里是求年龄的平均值，所以输入的数据是年龄
     * 由于是聚合函数，肯定时输入一个数组的数据，最后返回一个数据也就是平均值
     * 所以输入的是一个数组，数据的类别名叫age，数据的类型是longType
     */
    override def inputSchema: StructType = {
      StructType(
        Array(
          StructField("age", LongType)
        )
      )
    }

    /**
     * 缓冲区
     * 缓冲区是用来暂时存储数据，数据会在这里进行暂时的存储、运算然后才输出数据
     * 例如求平均值：数据在缓冲区进行求和和计算数量，求出平均值后输出
     *
     * @return
     */
    override def bufferSchema: StructType = {
      StructType(
        Array(
          StructField("total", LongType),
          StructField("count", LongType)
        )
      )
    }

    /**
     * 函数输出的数据类型就是是计算结果的数据类型
     *
     * @return
     */
    override def dataType: DataType = LongType

    /**
     * 函数的稳定性
     *
     * @return
     */
    override def deterministic: Boolean = true

    /**
     * 缓冲区的初始换
     *
     * @param buffer
     */
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      //这里就是如何该初始哈缓冲区的数据（也就是归零），这里有两个方法来归零
      //方法一
      //buffer(0) = 0l
      //buffer(1) = 0l

      //方法二
      buffer.update(0, 0l)
      buffer.update(1, 0l)
    }

    /**
     * 根据输入的数据来更新缓冲区的数据，也就是缓冲区的计算规则
     *
     * @param buffer
     * @param input
     */
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      //第一个数据就是求和，缓冲区里的数据加上输入的数据
      buffer.update(0, buffer.getLong(0) + input.getLong(0))
      //第二个数据就是计算总数，每次加一即可
      buffer.update(1, buffer.getLong(1) + 1)
    }

    /**
     * 缓冲区的数据合并
     * 保留1
     *
     * @param buffer1
     * @param buffer2
     */
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1.update(0, buffer1.getLong(0) + buffer2.getLong(0))
      buffer1.update(1, buffer1.getLong(1) + buffer2.getLong(1))
    }

    /**
     * 计算平均值
     *
     * @param buffer
     * @return
     */
    override def evaluate(buffer: Row): Any = (buffer.getLong(0) / buffer.getLong(1))
  }
}
